import jieba
import numpy as np
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.naive_bayes import MultinomialNB
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score

# 读取名字和性别数据
with open("NameData.txt", "r", encoding="latin-1") as file:
    data = [line.strip().split() for line in file]

# 将名字和性别分开
names, genders = zip(*data)

# 分词
seg_names = [' '.join(jieba.cut(name)) for name in names]

# 使用CountVectorizer将名字转换为特征向量
vectorizer = CountVectorizer()
X = vectorizer.fit_transform(seg_names)

# 将性别转换为数字标签（0表示男性，1表示女性）
y = np.array([0 if gender == '男' else 1 for gender in genders])

# 划分数据集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 创建并训练朴素贝叶斯分类器
classifier = MultinomialNB()
classifier.fit(X_train, y_train)

# 进行预测
y_pred = classifier.predict(X_test)

# 输出准确率
accuracy = accuracy_score(y_test, y_pred)
print(f"准确率: {accuracy}")

# 使用训练好的分类器进行性别预测
def predict_gender(name):
    seg_name = ' '.join(jieba.cut(name))
    name_vectorized = vectorizer.transform([seg_name])
    predicted_gender = classifier.predict(name_vectorized)
    return "男" if predicted_gender[0] == 0 else "女"

# 例子
test_name = "张三"
predicted_gender = predict_gender(test_name)
print(f"名字'{test_name}'的性别预测为: {predicted_gender}")
